import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# 加载数据
df = pd.read_csv('state_data_final.csv')

# 设置样式
plt.style.use('ggplot')
sns.set_palette("husl")

# 创建画布
plt.figure(figsize=(12, 8))

# 人口TOP10
plt.subplot(1, 2, 1)
top10 = df.nlargest(10, 'POPESTIMATE2019')
sns.barplot(x='POPESTIMATE2019', y='STATE', data=top10)
plt.title('人口最多的10个州')
plt.xlabel('人口数量')
plt.ylabel('')

# 人口BOTTOM10
plt.subplot(1, 2, 2)
bottom10 = df.nsmallest(10, 'POPESTIMATE2019')
sns.barplot(x='POPESTIMATE2019', y='STATE', data=bottom10)
plt.title('人口最少的10个州')
plt.xlabel('人口数量')
plt.ylabel('')

plt.tight_layout()
plt.savefig('population_top_bottom.png', dpi=300, bbox_inches='tight')
plt.close()
print("✅ 人口TOP10/BOTTOM10图表已保存为 population_top_bottom.png")
plt.figure(figsize=(10, 8))
sns.scatterplot(x='state_area', y='POPESTIMATE2019',
               size='density', hue='density',
               sizes=(50, 300), data=df)
plt.title('人口-面积-密度关系')
plt.xlabel('面积 (平方英里)')
plt.ylabel('人口数量')
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

plt.tight_layout()
plt.savefig('population_density_scatter.png', dpi=300, bbox_inches='tight')
plt.close()
print("✅ 人口密度散点图已保存为 population_density_scatter.png")
plt.figure(figsize=(12, 8))
density_pivot = df.pivot_table(index='STATE', values='density', aggfunc='mean').sort_values('density')
sns.heatmap(density_pivot, annot=True, fmt=".1f", cmap='YlOrRd')
plt.title('各州人口密度 (人/平方英里)')
plt.ylabel('')

plt.tight_layout()
plt.savefig('density_heatmap.png', dpi=300, bbox_inches='tight')
plt.close()
print("✅ 人口密度热力图已保存为 density_heatmap.png")
plt.figure(figsize=(12, 8))
plt.scatter(x=df['long'], y=df['lat'],
           s=df['POPESTIMATE2019']/100000,
           c=df['density'], cmap='viridis', alpha=0.7)
plt.colorbar(label='人口密度')
plt.title('人口地理分布 (气泡大小∝人口数量)')
plt.xlabel('经度')
plt.ylabel('纬度')

plt.tight_layout()
plt.savefig('geo_bubble.png', dpi=300, bbox_inches='tight')
plt.close()
print("✅ 地理分布气泡图已保存为 geo_bubble.png")
try:
    import folium

    print("\n正在生成交互式地图...")
    m = folium.Map(location=[37.8, -96], zoom_start=4)

    for idx, row in df.iterrows():
        folium.CircleMarker(
            location=[row['lat'], row['long']],
            radius=row['POPESTIMATE2019'] / 2000000,
            popup=f"{row['STATE']}<br>人口: {row['POPESTIMATE2019']:,}",
            color='blue',
            fill=True
        ).add_to(m)

    m.save('interactive_map.html')
    print("✅ 交互式地图已保存为 interactive_map.html")

except ImportError:
    print("ℹ️ 要生成交互式地图，请先安装folium: pip install folium")